import numpy as np
import h5py,os
import array
import scipy.io as sio

#### create all nodule as aim




def cr_data(main_dir):

    if not os.path.exists(main_dir+'/data/faster_rcnn_data/v9/JPEGImages'):
        os.makedirs(main_dir+'/data/faster_rcnn_data/v9/JPEGImages')
    if not os.path.exists(main_dir+'/data/faster_rcnn_data/v9/Annotations'):
        os.makedirs(main_dir+'/data/faster_rcnn_data/v9/Annotations')
    if not os.path.exists(main_dir+'/data/faster_rcnn_data/v9/ImageSets'):
        os.makedirs(main_dir+'/data/faster_rcnn_data/v9/ImageSets')
    if not os.path.exists(main_dir+'/data/faster_rcnn_data/v9/ImageSets/fold1'):
        os.makedirs(main_dir+'/data/faster_rcnn_data/v9/ImageSets/fold1')
    if not os.path.exists(main_dir+'/data/faster_rcnn_data/v9/ImageSets/fold2'):
        os.makedirs(main_dir+'/data/faster_rcnn_data/v9/ImageSets/fold2')
    if not os.path.exists(main_dir+'/data/faster_rcnn_data/v9/ImageSets/fold3'):
        os.makedirs(main_dir+'/data/faster_rcnn_data/v9/ImageSets/fold3')        
    if not os.path.exists(main_dir+'/data/faster_rcnn_data/v9/ImageSets/fold4'):
        os.makedirs(main_dir+'/data/faster_rcnn_data/v9/ImageSets/fold4')
    if not os.path.exists(main_dir+'/data/faster_rcnn_data/v9/ImageSets/fold5'):
        os.makedirs(main_dir+'/data/faster_rcnn_data/v9/ImageSets/fold5')

    f1=h5py.File(main_dir+'/data/lidc_v9/vol.hdf5','r')
    f2=h5py.File(main_dir+'/data/kaggle/vol.hdf5','r')
    l=np.load(main_dir+'/code/cr_frcnn_data/label_dict_v9.npy').item()
    fold_list = np.load(main_dir+'/code/cr_frcnn_data/fold_list.npy')
    # namelist_full=[]
    for fold_id in xrange(len(fold_list)):
        # namelist=[]
        for iscan in xrange(len(fold_list[fold_id])):
        
            name = fold_list[fold_id][iscan]
            print iscan, name

 			
            lab=l[name]
            if name[0]=='L':
                vol=f1[name][:]
            else:
            	vol=f2[name][:]
            

            vol = vol + 1350.0
            vol [vol < 0] = 0
            vol [vol > 1500] = 1500
            vol = vol * 255.0 /1500
            vol = np.floor(vol)
            vol [vol < 0] = 0   

            labf = [[] for _ in xrange(len(lab))]
            labl = [[] for _ in xrange(len(lab))]

            for ilab in xrange (len(lab)):
                h=lab[ilab][5]-lab[ilab][2]
                labf[ilab]=np.round(lab[ilab][2]+0.2*h)
                labl[ilab]=np.round(lab[ilab][2]+0.8*h)

            for islice in xrange(1,vol.shape[2]-1):
                cr_img = 0
                cr_xml = 0

                for j in xrange (len(lab)):
                    if (islice >=labf[j]) and (islice<=labl[j]):       
                        if cr_img == 0:
                            img=np.dstack((vol[:,:,islice-1],vol[:,:,islice],vol[:,:,islice+1]))
                            sio.savemat(main_dir+'/data/faster_rcnn_data/v9/JPEGImages/'+name+'_'+str(islice)+'.mat',{'img':img},do_compression=True)
                            # namelist.append(name+'_'+str(islice))
                            if cmp(name,'LIDC-IDRI-0671_01')==0 and islice==50:
                                sio.savemat(main_dir+'/data/faster_rcnn_data/v9/JPEGImages/'+name+'_'+str(islice)+'.mat',{'img':img})
                            
                            cr_img=1

                        if cr_xml==0:
                            fid=open(main_dir+'/data/faster_rcnn_data/v9/Annotations/'+name+'_'+str(islice)+'.xml','w')
                            print >> fid,'<annotation>\n'
                            print >> fid,'\t<size>\n'
                            print >>fid,'\t\t<width>512</width>\n'
                            print >> fid,'\t\t<height>512</height>\n'
                            print >> fid,'\t\t<depth>3</depth>\n'
                            print >>fid,'\t</size>\n'
                            print >>fid,'\t<object>\n'
                            print >>fid,'\t\t<name>nodule</name>\n'
                            print >>fid,'\t\t<difficult>0</difficult>\n'
                            print >>fid,'\t\t<bndbox>\n'
                            print >>fid,'\t\t\t<xmin>', int(np.round(lab[j][1])), '</xmin>\n'
                            print >>fid,'\t\t\t<ymin>', int(np.round(lab[j][0])), '</ymin>\n'
                            print >>fid,'\t\t\t<xmax>', int(np.round(lab[j][4])), '</xmax>\n'
                            print >>fid,'\t\t\t<ymax>', int(np.round(lab[j][3])), '</ymax>\n'
                            print>>fid,'\t\t</bndbox>\n'
                            print>>fid,'\t</object>\n'
                            cr_xml=1
                        else:
                            print >>fid,'\t<object>\n'
                            print >>fid,'\t\t<name>nodule</name>\n'
                            print >>fid,'\t\t<difficult>0</difficult>\n'
                            print >>fid,'\t\t<bndbox>\n'
                            print >>fid,'\t\t\t<xmin>', int(np.round(lab[j][1])), '</xmin>\n'
                            print >>fid,'\t\t\t<ymin>', int(np.round(lab[j][0])), '</ymin>\n'
                            print >>fid,'\t\t\t<xmax>', int(np.round(lab[j][4])), '</xmax>\n'
                            print >>fid,'\t\t\t<ymax>', int(np.round(lab[j][3])), '</ymax>\n'
                            print >>fid,'\t\t</bndbox>\n'
                            print >>fid,'\t</object>\n'
                if cr_xml:
                    print >>fid,'</annotation>\n'
                    fid.close()
        # namelist_full.append(namelist)

    # for mm in xrange(5):
    #     fid=open(main_dir+'/data/faster_rcnn_data/v9/ImageSets/fold'+str(mm+1)+'/trainval.txt','w')
        
    #     for nn in xrange(5):
    #         if nn!=mm: 
    #             for ilist in xrange(len(namelist_full[nn])):
    #                 print >>fid,namelist_full[nn][ilist]
    #     fid.close() 

def cr_list(main_dir):
    fold_list = np.load(main_dir+'/code/cr_frcnn_data/fold_list.npy')
    D=os.listdir(main_dir+'/data/faster_rcnn_data/v9/Annotations/')
    D=sorted(D,key=str.lower)
    new_list=[]
    for i in xrange(len(D)):
        name=D[i]
        if name[0]=='L':
            new_list.append(name[0:17])
        else:
            new_list.append(name[0:32])
    
    
    namelist_full=[]
    for ifold in xrange(5):
        name_list=[]
        for iscan in xrange(len(fold_list[ifold])):
            foldname=fold_list[ifold][iscan]
            for j in xrange(len(new_list)):
                if cmp(foldname,new_list[j])==0:
                    name_list.append(D[j][0:len(D[j])-4])
  
        namelist_full.append(name_list)

    for mm in xrange(5):
        fid=open(main_dir+'/data/faster_rcnn_data/v9/ImageSets/fold'+str(mm+1)+'/trainval.txt','w')
        
        for nn in xrange(5):
            if nn!=mm: 
                for ilist in xrange(len(namelist_full[nn])):
                    print >>fid,namelist_full[nn][ilist]
        fid.close() 






if __name__ == '__main__':
    #main_dir='/data_4t/Kaggle/submission_code_data'
    main_dir = os.getcwd() + '/../..'
    cr_data(main_dir)
    cr_list(main_dir)
